#include "QuadTree.h"

bool Rect::intersect(const Segment& seg) const {
    double minX = min(seg.a.x, seg.b.x);
    double maxX = max(seg.a.x, seg.b.x);
    double minY = min(seg.a.y, seg.b.y);
    double maxY = max(seg.a.y, seg.b.y);
    return !(maxX < x0 || minX > x1 || maxY < y0 || minY > y1);
}

QuadNode::QuadNode(Rect r) : rect(r), isLeaf(true) {}

QuadNode::~QuadNode() {
    for (auto child : children) {
        delete child;
    }
}

QuadTree::QuadTree() : root(new QuadNode(Rect{ 0, 0, 65535.0, 65535.0 })), maxDepth(8), maxSegsPerLeaf(40) {}

QuadTree::~QuadTree() {
    delete root;
}

void QuadTree::split(QuadNode* node, int depth) {
    double x0 = node->rect.x0, y0 = node->rect.y0;
    double x1 = node->rect.x1, y1 = node->rect.y1;
    double midx = x0 + (x1 - x0) / 2;
    double midy = y0 + (y1 - y0) / 2;

    node->children = {
        new QuadNode(Rect{x0, y0, midx, midy}),
        new QuadNode(Rect{midx, y0, x1, midy}),
        new QuadNode(Rect{x0, midy, midx, y1}),
        new QuadNode(Rect{midx, midy, x1, y1})
    };
    node->isLeaf = false;

    vector<int> segsToKeep;
    for (int segIndex : node->segs) {
        Segment& s = segments[segIndex];
        double minx = min(s.a.x, s.b.x);
        double maxx = max(s.a.x, s.b.x);
        double miny = min(s.a.y, s.b.y);
        double maxy = max(s.a.y, s.b.y);

        if (minx <= node->rect.x0 && maxx >= node->rect.x1 &&
            miny <= node->rect.y0 && maxy >= node->rect.y1) {
            segsToKeep.push_back(segIndex);
        }
        else {
            for (auto child : node->children) {
                if (child->rect.intersect(s)) {
                    insertIntoNode(segIndex, depth + 1, child);
                }
            }
            segToNodes[segIndex].erase(node);
        }
    }
    node->segs = segsToKeep;
}

void QuadTree::insertIntoNode(int segIndex, int depth, QuadNode* node) {
    if (node->isLeaf) {
        if (node->segs.size() >= maxSegsPerLeaf && depth < maxDepth) {
            split(node, depth);
            insertIntoNode(segIndex, depth, node);
            return;
        }
        else {
            node->segs.push_back(segIndex);
            segToNodes[segIndex].insert(node);
        }
    }
    else {
        Segment& s = segments[segIndex];
        double minx = min(s.a.x, s.b.x);
        double maxx = max(s.a.x, s.b.x);
        double miny = min(s.a.y, s.b.y);
        double maxy = max(s.a.y, s.b.y);

        const double eps = 1e-6;
        if (minx <= node->rect.x0 + eps &&
            maxx >= node->rect.x1 - eps &&
            miny <= node->rect.y0 + eps &&
            maxy >= node->rect.y1 - eps) {
            node->segs.push_back(segIndex);
            segToNodes[segIndex].insert(node);
        }
        else {
            for (auto child : node->children) {
                if (child->rect.intersect(s)) {
                    insertIntoNode(segIndex, depth + 1, child);
                }
            }
        }
    }
}

double QuadTree::calcDistRectSq(double x, double y, const Rect& rect) const {
    double dx = 0, dy = 0;
    if (x < rect.x0) dx = rect.x0 - x;
    else if (x > rect.x1) dx = x - rect.x1;
    if (y < rect.y0) dy = rect.y0 - y;
    else if (y > rect.y1) dy = y - rect.y1;
    return dx * dx + dy * dy;
}

double QuadTree::calcDistSq(double x, double y, const Segment& seg) const {
    const Point& A = seg.a;
    double apx = x - A.x;
    double apy = y - A.y;

    const double eps = 1e-9;
    if (seg.len2 < eps) {
        return apx * apx + apy * apy;
    }

    double dot = apx * seg.dx + apy * seg.dy;
    double t = dot / seg.len2;

    if (t < 0.0) {
        return apx * apx + apy * apy;
    }
    if (t > 1.0) {
        double bpx = x - seg.b.x;
        double bpy = y - seg.b.y;
        return bpx * bpx + bpy * bpy;
    }

    double projx = A.x + t * seg.dx;
    double projy = A.y + t * seg.dy;
    double dx = x - projx;
    double dy = y - projy;
    return dx * dx + dy * dy;
}

int QuadTree::addSegment(double x1, double y1, double x2, double y2, vector<vector<wstring>>& queryPointData) {
    Segment seg = { Point{x1, y1}, Point{x2, y2},
             x2 - x1, y2 - y1,
             (x2 - x1) * (x2 - x1) + (y2 - y1) * (y2 - y1) };
    segments.push_back(seg);
    int segIndex = segments.size() - 1;
    insertIntoNode(segIndex, 0, root);

#pragma omp parallel
    {
        vector<pair<int, double>> localUpdates;

#pragma omp for
        for (size_t queryIdx = 0; queryIdx < queryPoints.size(); ++queryIdx) {
            auto xy = queryPoints[queryIdx];
            double newDistSq = calcDistSq(xy.first, xy.second, seg);
            auto& currentResults = queryResults[queryIdx];

            if (currentResults.empty() || newDistSq < currentResults.back().second * currentResults.back().second || currentResults.back().first == 0) {
                localUpdates.emplace_back(queryIdx, sqrt(newDistSq));
            }
        }

#pragma omp critical
        for (auto& update : localUpdates) {
            int qIdx = update.first;
            double dist = update.second;
            for (auto& oldRes : queryResults[qIdx]) {
                segToQueries[oldRes.first].erase(qIdx);
            }
            queryResults[qIdx].clear();
            queryResults[qIdx].emplace_back(segIndex + 1, dist);
            queryPointData[qIdx + 1][1] = to_wstring(segIndex + 1);
            queryPointData[qIdx + 1][2] = to_wstring(dist);
            segToQueries[segIndex + 1].insert(qIdx);
        }
    }

    return segIndex + 1;
}

void QuadTree::removeSegment(int segId, vector<vector<wstring>>& queryPointData) {
    int segIndex = segId - 1;
    if (segIndex < 0 || static_cast<size_t>(segIndex) >= segments.size()) return;

    auto affectedQueries = segToQueries[segId];
    segToQueries.erase(segId);

    for (auto node : segToNodes[segIndex]) {
        auto it = find(node->segs.begin(), node->segs.end(), segIndex);
        if (it != node->segs.end()) {
            node->segs.erase(it);
        }
    }
    segToNodes.erase(segIndex);

#pragma omp parallel for
    for (int queryIdx : affectedQueries) {
        auto xy = queryPoints[queryIdx];
        auto newRes = queryInternal(xy.first, xy.second, queryIdx);
        queryPointData[queryIdx + 1][1] = to_wstring(newRes.first);
        queryPointData[queryIdx + 1][2] = to_wstring(newRes.second);
    }
}

pair<int, double> QuadTree::queryInternal(double x, double y, int queryIdx) {
    double minDistSq = DBL_MAX;
    int bestSegId = -1;

    auto cmp = [](const pair<double, QuadNode*>& a, const pair<double, QuadNode*>& b) {
        return a.first > b.first;
        };
    priority_queue<pair<double, QuadNode*>, vector<pair<double, QuadNode*>>, decltype(cmp)> pq(cmp);
    pq.push({ 0, root });

    while (!pq.empty()) {
        double dBound = pq.top().first;
        QuadNode* node = pq.top().second;
        pq.pop();

        if (dBound >= minDistSq) continue;

#pragma omp parallel
        {
            double localMin = DBL_MAX;
            int localId = -1;

#pragma omp for
            for (int segIndex : node->segs) {
                double dSq = calcDistSq(x, y, segments[segIndex]);
                if (dSq < localMin ||
                    (abs(dSq - localMin) < 1e-9 && segIndex < localId)) {
                    localMin = dSq;
                    localId = segIndex;
                }
            }

#pragma omp critical
            {
                if (localMin < minDistSq ||
                    (abs(localMin - minDistSq) < 1e-9 && localId < bestSegId)) {
                    minDistSq = localMin;
                    bestSegId = localId;
                }
            }
        }

        if (minDistSq == 0) break;

        if (!node->isLeaf) {
#pragma omp parallel for schedule(dynamic)
            for (auto child : node->children) {
                double dChild = calcDistRectSq(x, y, child->rect);
                if (dChild < minDistSq) {
#pragma omp critical
                    pq.push({ dChild, child });
                }
            }
        }
    }
    minDistSq = (minDistSq == DBL_MAX) ? -1 : sqrt(minDistSq);
    pair<int, double> result(bestSegId + 1, minDistSq);

    if (queryIdx != -1) {
#pragma omp critical
        {
            queryResults[queryIdx].push_back(result);
            segToQueries[result.first].insert(queryIdx);
        }
    }

    return result;
}

pair<int, double> QuadTree::query(double x, double y) {
    int queryIdx = queryPoints.size();
    queryPoints.emplace_back(x, y);
    return queryInternal(x, y, queryIdx);
}

void QuadTree::clear() {
    delete root;
    root = new QuadNode(Rect{ 0.0, 0.0, 65535.0, 65535.0 });
    segments.clear();
    queryPoints.clear();
    queryResults.clear();
    segToNodes.clear();
    segToQueries.clear();
}